#include <strmif.h>
#include "mkvparserstreamvideo.hpp"
#include "mkvparser.hpp"
#include <cassert>
//#include <algorithm>
//#include <limits>
//#include <climits>
//#include <vfwmsgs.h>
//#include <sstream>
//#include <iomanip>
#include <dvdmedia.h>
#include <uuids.h>
//#include <mmreg.h>
#ifdef _DEBUG
#include "odbgstream.hpp"
using std::endl;
#endif

namespace MkvParser
{


VideoStream* VideoStream::CreateInstance(VideoTrack* pTrack)
{
    assert(pTrack);
    
    const char* const id = pTrack->GetCodecId();
    assert(id);  //TODO
    
    if (_stricmp(id, "V_MPEG2") == 0)
        __noop;
    else
        return 0;  //can't create a stream from this track
        
    //TODO: vet settings, etc
    
    //At least one cluster should have been loaded when we opened
    //the file.  We search the first cluster for a frame having
    //this track number, and use that to start from.  If no frame
    //with this track number is present in first cluster, then
    //we assume (correctly or incorrectly) that this file doesn't
    //have any frames from that track at all.
    
    VideoStream* const s = new (std::nothrow) VideoStream(pTrack);
    assert(s);  //TODO
    
    return s;
}


VideoStream::VideoStream(VideoTrack* pTrack) : Stream(pTrack)
{
}


std::wostream& VideoStream::GetKind(std::wostream& os) const
{
    return os << L"Video";
}


void VideoStream::GetMediaTypes(CMediaTypes& mtv) const
{
    mtv.Clear();
    
    const char* const id = m_pTrack->GetCodecId();
    assert(id);
    
    if (_stricmp(id, "V_MPEG2") == 0)
        GetMpeg2MediaTypes(mtv);        
}


void VideoStream::GetMpeg2MediaTypes(CMediaTypes& mtv) const
{
    const bytes_t& cp = m_pTrack->GetCodecPrivate();
    const ULONG cp_size = static_cast<ULONG>(cp.size());
    assert(cp_size >= 4);
    assert(cp[0] == 0);
    assert(cp[1] == 0);
    assert(cp[2] == 1);
    assert(cp[3] == 0xB3);
    
    AM_MEDIA_TYPE mt;
    
    //TODO: use offsetof instead?
    const ULONG off = FIELD_OFFSET(MPEG2VIDEOINFO, dwSequenceHeader[0]);
    const ULONG cbFormat = off + cp_size;

    BYTE* const pbFormat = (BYTE*)_alloca(cbFormat);
    MPEG2VIDEOINFO& mvi = (MPEG2VIDEOINFO&)(*pbFormat);
    
    VIDEOINFOHEADER2& vih = mvi.hdr;
    BITMAPINFOHEADER& bmih = vih.bmiHeader;

    mt.majortype = MEDIATYPE_Video;
    mt.subtype = MEDIASUBTYPE_MPEG2_VIDEO;
    mt.bFixedSizeSamples = FALSE;
    mt.bTemporalCompression = TRUE;
    mt.lSampleSize = 0;
    mt.formattype = FORMAT_MPEG2_VIDEO;
    mt.pUnk = 0;
    mt.cbFormat = cbFormat;
    mt.pbFormat = pbFormat;
    
    SetRectEmpty(&vih.rcSource);  //TODO
    SetRectEmpty(&vih.rcTarget);
    vih.dwBitRate = 0;
    vih.dwBitErrorRate = 0;
    
    const VideoTrack* const pTrack = static_cast<VideoTrack*>(m_pTrack);
    
    const float r = pTrack->GetFrameRate();
    
    if (r <= 0)
        vih.AvgTimePerFrame = 0;
    else    
    {
        const float tt = 10000000 / r;  //[ticks/sec] / [frames/sec]
        vih.AvgTimePerFrame = static_cast<__int64>(tt);
    }
    
    vih.dwInterlaceFlags = 0;   //TODO
    vih.dwCopyProtectFlags = 0;
    vih.dwPictAspectRatioX = 0;  //TODO
    vih.dwPictAspectRatioY = 0;  //TODO
    vih.dwControlFlags = 0;      //TODO
    vih.dwReserved2 = 0;
    
    const __int64 w = pTrack->GetWidth();
    assert(w > 0);
    assert(w <= LONG_MAX);
    
    const __int64 h = pTrack->GetHeight();
    assert(h > 0);
    assert(h <= LONG_MAX);
    
    bmih.biSize = sizeof bmih;
    bmih.biWidth = static_cast<LONG>(w);
    bmih.biHeight = static_cast<LONG>(h);
    bmih.biPlanes = 1;  //TODO
    bmih.biBitCount = 0;  //TODO
    bmih.biCompression = mt.subtype.Data1;  //TODO
    bmih.biSizeImage = 0;
    bmih.biXPelsPerMeter = 0;
    bmih.biYPelsPerMeter = 0;
    bmih.biClrUsed = 0;
    bmih.biClrImportant = 0;
    
    mvi.dwStartTimeCode = 0;  //TODO
    mvi.cbSequenceHeader = cp_size;
    mvi.dwProfile = 0;  //TODO
    mvi.dwLevel = 0;   //TODO
    mvi.dwFlags = 0;  //TODO
    memcpy(pbFormat + off, &cp[0], cp.size());
    
    mtv.Add(mt);
}


HRESULT VideoStream::QueryAccept(const AM_MEDIA_TYPE* pmt) const
{
    if (pmt == 0)
        return E_INVALIDARG;
        
    const AM_MEDIA_TYPE& mt = *pmt;
    
    if (mt.majortype != MEDIATYPE_Video)
        return S_FALSE;
        
    const char* const id = m_pTrack->GetCodecId();
    assert(id);
    
    if (_stricmp(id, "V_MPEG2") == 0)
    {
        if (mt.subtype != MEDIASUBTYPE_MPEG2_VIDEO)
            return S_FALSE;
            
        return S_OK;  //TODO: more vetting here
    }

    return S_FALSE;  
}


HRESULT VideoStream::UpdateAllocatorProperties(
    ALLOCATOR_PROPERTIES& props) const
{
    if (props.cBuffers <= 0)
        props.cBuffers = 1;
        
    const long size = GetBufferSize();
    
    if (props.cbBuffer < size)
        props.cbBuffer = size;
        
    if (props.cbAlign <= 0)
        props.cbAlign = 1;
        
    if (props.cbPrefix < 0)
        props.cbPrefix = 0;
    
    return S_OK;
}


long VideoStream::GetBufferSize() const
{   
    const VideoTrack* const pTrack = static_cast<VideoTrack*>(m_pTrack);
        
    const __int64 w = pTrack->GetWidth();    
    const __int64 h = pTrack->GetHeight();

    const __int64 size_ = w * h * 4;  //RGB32 (worst case)
    assert(size_ <= LONG_MAX);
    
    const long size = static_cast<LONG>(size_);
    
    return size;
}


#if 0  //TODO: use this for incremental loading
HRESULT VideoStream::PopulateSample(IMediaSample* pSample, __int64& result)
{
    if (pSample == 0)
        return E_INVALIDARG;
        
    if (m_pCurr == m_pStop)
    {
#ifdef _DEBUG
        odbgstream os;

        os << "VideoStream::PopulateSample: EOS detected";
        
        if (m_pCurr)
            os << "; time=" << m_pCurr->GetTime();
            
        os << endl;
#endif
        
        result = 0;
        return S_FALSE;  //end-of-stream
    }

    assert(m_baseTime >= 0);
    assert((m_baseTime % 100) == 0);
    assert(m_pCurr);    
    assert(m_pCurr->GetNumber() == m_pTrack->GetNumber());
    //assert(m_pCurr->GetTime() >= m_baseTime);    
    assert((m_pStop == 0) ||
           (m_pStop->GetTimeCode() > m_pCurr->GetTimeCode()));
    
    const Block* const pNextBlock = m_pTrack->GetNext(m_pCurr, result);

    if (result < 0)
        return static_cast<HRESULT>(result);
        
    if (result > 0)
        return S_FALSE;  //can't make progress until we have next frame too
        
    const LONG srcsize = m_pCurr->GetSize();
    assert(srcsize >= 0);
    
    const long tgtsize = pSample->GetSize();
    tgtsize;
    assert(tgtsize >= 0);
    assert(tgtsize >= srcsize);
    
    BYTE* ptr;

    HRESULT hr = pSample->GetPointer(&ptr);  //read srcsize bytes
    assert(SUCCEEDED(hr));
    assert(ptr);
    
    hr = m_pCurr->Read(ptr);
    assert(hr == S_OK);  //all bytes were read
    
    hr = pSample->SetActualDataLength(srcsize);
    
    hr = pSample->SetPreroll(0);
    assert(SUCCEEDED(hr));
    
    hr = pSample->SetMediaType(0);
    assert(SUCCEEDED(hr));
    
    hr = pSample->SetDiscontinuity(m_bDiscontinuity);
    assert(SUCCEEDED(hr));
    m_bDiscontinuity = false;
    
    hr = pSample->SetMediaTime(0, 0);
    assert(SUCCEEDED(hr));
    
    const BlockGroup* const pGroup = m_pCurr->m_pGroup;
    assert(pGroup);
    
    const bool bKey = (pGroup->GetPrevTimeCode() >= 0);
    
    hr = pSample->SetSyncPoint(bKey ? TRUE : FALSE);
    assert(SUCCEEDED(hr));
    
    if (pGroup->GetNextTimeCode() > 0)  //B-frame
    {
        assert(!bKey);
        
        hr = pSample->SetTime(0, 0);
        assert(SUCCEEDED(hr));
    }
    else
    {
        const __int64 start_ns = m_pCurr->GetTime();
        assert(start_ns >= m_baseTime);
        assert((start_ns % 100) == 0);
        
        const __int64 ns = start_ns - m_baseTime;        
        __int64 start_reftime = ns / 100;
#if 0
        //TODO: we can't do this, because frame don't appear in timecode order.
        //One possibility is to search for the next frame that isn't a B-frame.
                
        __int64 stop_reftime;
        __int64* pstop_reftime;
                
        if (pNextBlock == 0)
            pstop_reftime = 0;  //TODO: use duration of curr block
        else
        {
            const __int64 stop_tc = pNextBlock->GetTimeCode();
            assert(stop_tc > start_tc);
            
            const __int64 stop_ns = scale * stop_tc;
            assert((stop_ns % 100) == 0);

            stop_reftime = stop_ns / 100;
            assert(stop_reftime > start_reftime);

            pstop_reftime = &stop_reftime;
        }
        
        hr = pSample->SetTime(&start_reftime, pstop_reftime);
        assert(SUCCEEDED(hr));
#else
        hr = pSample->SetTime(&start_reftime, 0);
        assert(SUCCEEDED(hr));
#endif
    }
    
    m_pCurr = pNextBlock;    
    return S_OK;
}
#endif


HRESULT VideoStream::PopulateSample(IMediaSample* pSample)
{
    if (pSample == 0)
        return E_INVALIDARG;
        
    if (m_pCurr == m_pStop)
    {
#ifdef _DEBUG
        odbgstream os;

        os << "VideoStream::PopulateSample: EOS detected";
        
        if (m_pCurr)
            os << "; time=" << m_pCurr->GetTime();
            
        os << endl;
#endif
        
        return S_FALSE;  //end-of-stream
    }

    assert(m_baseTime >= 0);
    assert((m_baseTime % 100) == 0);
    assert(m_pCurr);    
    assert(m_pCurr->GetNumber() == m_pTrack->GetNumber());
    //assert(m_pCurr->GetTime() >= m_baseTime);    
    assert((m_pStop == 0) ||
           (m_pStop->GetTimeCode() > m_pCurr->GetTimeCode()));
    
    const Block* const pNextBlock = m_pTrack->GetNext(m_pCurr);

    const LONG srcsize = m_pCurr->GetSize();
    assert(srcsize >= 0);
    
    const long tgtsize = pSample->GetSize();
    tgtsize;
    assert(tgtsize >= 0);
    assert(tgtsize >= srcsize);
    
    BYTE* ptr;

    HRESULT hr = pSample->GetPointer(&ptr);  //read srcsize bytes
    assert(SUCCEEDED(hr));
    assert(ptr);
    
    hr = m_pCurr->Read(ptr);
    assert(hr == S_OK);  //all bytes were read
    
    hr = pSample->SetActualDataLength(srcsize);
    
    hr = pSample->SetPreroll(0);
    assert(SUCCEEDED(hr));
    
    hr = pSample->SetMediaType(0);
    assert(SUCCEEDED(hr));
    
    hr = pSample->SetDiscontinuity(m_bDiscontinuity);
    assert(SUCCEEDED(hr));
    m_bDiscontinuity = false;
    
    hr = pSample->SetMediaTime(0, 0);
    assert(SUCCEEDED(hr));
    
    const BlockGroup* const pGroup = m_pCurr->m_pGroup;
    assert(pGroup);
    
    const bool bKey = (pGroup->GetPrevTimeCode() >= 0);
    
    hr = pSample->SetSyncPoint(bKey ? TRUE : FALSE);
    assert(SUCCEEDED(hr));
    
    if (pGroup->GetNextTimeCode() > 0)  //B-frame
    {
        assert(!bKey);
        
        hr = pSample->SetTime(0, 0);
        assert(SUCCEEDED(hr));
    }
    else
    {
        const __int64 start_ns = m_pCurr->GetTime();
        assert(start_ns >= m_baseTime);
        assert((start_ns % 100) == 0);
        
        const __int64 ns = start_ns - m_baseTime;        
        __int64 start_reftime = ns / 100;
#if 0
        //TODO: we can't do this, because frame don't appear in timecode order.
        //One possibility is to search for the next frame that isn't a B-frame.
                
        __int64 stop_reftime;
        __int64* pstop_reftime;
                
        if (pNextBlock == 0)
            pstop_reftime = 0;  //TODO: use duration of curr block
        else
        {
            const __int64 stop_tc = pNextBlock->GetTimeCode();
            assert(stop_tc > start_tc);
            
            const __int64 stop_ns = scale * stop_tc;
            assert((stop_ns % 100) == 0);

            stop_reftime = stop_ns / 100;
            assert(stop_reftime > start_reftime);

            pstop_reftime = &stop_reftime;
        }
        
        hr = pSample->SetTime(&start_reftime, pstop_reftime);
        assert(SUCCEEDED(hr));
#else
        hr = pSample->SetTime(&start_reftime, 0);
        assert(SUCCEEDED(hr));
#endif
    }
    
    m_pCurr = pNextBlock;    
    return S_OK;
}




}  //end namespace MkvParser
